import os
import shutil
import glob
import pandas as pd
import numpy as np
from PIL import Image
import h5py


category2id = {
    # 'Reference': '1',
    # 'Surface_particle': '7',
    # 'CMP_scratch': '6',
    # 'Burried_particle': '4',
    # 'Void': '107',
    # 'Residue': '11',
    # 'Bump': '3',
    # 'Oxide_scratch': '5',
    # 'ADI_particle_deve': '67',
    # 'Pre_layer_def': '14',
    # 'Cu_loss': '84',
    # 'Cu_diffuse': '83',
    # 'Pattern_fail': '74',
    # 'ArrayPeeling': '88',
    # 'OtherPeeling': '89'


    #数据集种类扩充后
    'Reference': '1',
    'Surface_particle': '7',
    'Void': '107',
    'Residue': '11',
    'Burried_particle': '4',
    'Bump': '3',
    'CMP_scratch': '6',
    'Oxide_scratch': '5',
    'Prelayer_defect_developed': '14',
    'Cu_missing': '84',
    'Cu_diffuse': '83',
    'ADI_particle_developed': '67',
    'Pattern_fail': '74',
    'Other_peeling': '89',
    'Array_peeling': '88',
     #新增
    # 'EG_pattern_fail': '174',
    'Flake': '24',
    'Partial_etch': '76',
    'W_residue': '13',
    # 'Over_polish': '35',
    # 'LITHO_defocus': '52',
    # 'Corrosion': '62',
    # 'Metal_grain': '71',
    'Bridge': '12',
    # 'Ball_type': '21',
    'OX_residue': '27',
    'Seam': '33',
    # 'EG_other_peeling': '189',
    # 'EG_prelayer_defect_developed': '114',
    'Met_bridge': '214',
    'Met_residue': '68',
    'PR_peeling': '19',
    # 'SLICE_particle': '65',
    # 'Tiny_particle': '8'
}

id2category = {
    # '1': 'Reference',
    # '7': 'Surface_particle',
    # '6': 'CMP_scratch',
    # '4': 'Burried_particle',
    # '107': 'Void',
    # '11': 'Residue',
    # '3': 'Bump',
    # '5': 'Oxide_scratch',
    # '67': 'ADI_particle_deve',
    # '14': 'Pre_layer_def',
    # '84': 'Cu_loss',
    # '83': 'Cu_diffuse',
    # '74': 'Pattern_fail',
    # '88': 'ArrayPeeling',
    # '89': 'OtherPeeling',


    #数据集种类扩充后

    '1': 'Reference',
    '7': 'Surface_particle',
    '107': 'Void',
    '11': 'Residue',
    '4': 'Burried_particle',
    '3': 'Bump',
    '6': 'CMP_scratch',
    '5': 'Oxide_scratch',
    '14': 'Prelayer_defect_developed',
    '84': 'Cu_missing',
    '83': 'Cu_diffuse',
    '67': 'ADI_particle_developed',
    '74': 'Pattern_fail',
    '89': 'Other_peeling',
    '88': 'Array_peeling',
     #新增
    # '174': 'EG_pattern_fail',
    '24': 'Flake',
    '76': 'Partial_etch',
    '13': 'W_residue',
    # '35': 'Over_polish',
    # '52': 'LITHO_defocus',
    # '62': 'Corrosion',
    # '71': 'Metal_grain',
    '12': 'Bridge',
    # '21': 'Ball_type',
    '27': 'OX_residue',
    '33': 'Seam',
    # '189': 'EG_other_peeling',
    # '114': 'EG_prelayer_defect_developed',
    '214': 'Met_bridge',
    '68': 'Met_residue',
    '19': 'PR_peeling',
    # '65': 'SLICE_particle',
    # '8': 'Tiny_particle'
}


def readKlarf(filePath):
    dataAreaFlag = False
    defectDict = {}
    fileRowData = []
    with open(filePath, 'r') as f:
        for rowData in f.readlines():
            fileRowData.append(rowData)
            if rowData.startswith('SummarySpec'):
                dataAreaFlag = False
            if dataAreaFlag:
                rowList = rowData.split(' ')
                defectDict[int(rowList[1])] = int(rowList[10])
            if rowData.startswith('DefectList'):
                dataAreaFlag = True
        f.close()

    df = pd.read_csv(filePath, names=['col1', 'col2'], header=None, sep=" ")
    deviceId = str(df.loc['DeviceID', 'col1'])[0:-1]
    stepId = str(df.loc['StepID', 'col1'])[0:-1]
    lotId = 'l' + str(df.loc['LotID', 'col1'])[0:-1]
    waferId = 'w' + str(df.loc['WaferID', 'col1'])[0:-1]
    return deviceId, stepId, lotId, waferId, defectDict



def check_data():
    dir = r'C:\Users\caoxijian\Desktop\origin_datasets\klarf\total231218_sample'
    klarf_files = os.listdir(dir)
    print(len(klarf_files))

    total_defect = 0
    fileDirs = set()
    defect_type = {}
    for file in klarf_files:
        filePath = os.path.join(dir, file)
        deviceId, stepId, lotId, waferId, defectDict = readKlarf(filePath)
        fileDir = deviceId + '__' + stepId + '__' + lotId + '__' + waferId
        fileDirs.add(fileDir)
        total_defect = total_defect + len(defectDict)
        # print(f"len of defect: {len(defectDict)}")
        for v in defectDict.values():
            if v in defect_type.keys():
                defect_type[v] += 1
            else:
                defect_type[v] = 1

    categories2ids = {}
    fileDirs = sorted(fileDirs)
    print(fileDirs)
    print(len(fileDirs))
    print(total_defect)
    print(defect_type)
    defect_type = sorted(defect_type.items(), key=lambda x: x[1], reverse=True)
    print(defect_type)
    print(type(defect_type))
    print(len(defect_type))

def sampleFromklarf():
    src_dir = r'C:\Users\caoxijian\Desktop\origin_datasets\klarf\total231218'
    dst_dir = r'C:\Users\caoxijian\Desktop\origin_datasets\klarf\total231218_sample'
    if os.path.exists(dst_dir):
        shutil.rmtree(dst_dir)
    os.mkdir(dst_dir)
    klarf_list = os.listdir(src_dir)
    np.random.shuffle(klarf_list)
    klarf_sample_list = klarf_list[0:3500]
    for klarf in klarf_sample_list:
        src_klarf = os.path.join(src_dir, klarf)
        dst_klarf = os.path.join(dst_dir, klarf)
        shutil.copyfile(src_klarf, dst_klarf)



def getH5Data():
    dir = r'C:\Users\caoxijian\Desktop\origin_datasets\klarf\total231218'
    dst_dir = r'D:\proj\dataset_h5_2'

    if os.path.exists(dst_dir):
        shutil.rmtree(dst_dir)
    os.mkdir(dst_dir)

    for category in category2id.keys():
        category_dir = os.path.join(dst_dir, category)
        if not os.path.exists(category_dir):
            os.mkdir(category_dir)

    klarf_files = os.listdir(dir)
    for file in klarf_files:
        filePath = os.path.join(dir, file)
        deviceId, stepId, lotId, waferId, defectDict = readKlarf(filePath)
        src_dir = os.path.join('Z:\\' + deviceId, stepId, lotId, waferId)
        for k, v in defectDict.items():  # k: defect_id   v: defect_type
            if str(v) in category2id.values():
                category = id2category[str(v)]
                dst_h5_file_name = category + '_' + deviceId + '_' + stepId + '_' + lotId + '_' + waferId + '_' + str(
                        k) + '_' + 'Topography3.h5'
                dst_h5_file = os.path.join(dst_dir, category, dst_h5_file_name)
                f = h5py.File(dst_h5_file, 'w')
                img_dir = src_dir + r'\*\Images\Run_1\*'
                img_list = glob.glob((img_dir))
                # ['Z:\\Y-MS42\\AL_ASI_MINI\\lDE1621.04_0718\\w21\\BB94982A-9250-4CF7-9B04-71E82FCB2B87\\Images\\Run_1\\Defect_99982_Class_1_Topography3.tiff']
                for img in img_list:
                    if 'class_tiff' in f.keys() and 'defect_tiff' in f.keys() and 'reference_tiff' in f.keys():
                        break
                    if 'Defect_' + str(k) + '_Class_1_Topography3' in img:
                        if 'class_tiff' not in f.keys():
                            class_tiff = np.array(Image.open(img))
                            print(f"add class_file from {img} to {dst_h5_file}")
                            f.create_dataset('class_tiff', data=class_tiff)

                    elif 'Defect_' + str(k) + '_Defect_1_Topography3' in img:
                        if 'defect_tiff' not in f.keys():
                            defect_tiff = np.array(Image.open(img))
                            print(f"add defect_file from {img} to {dst_h5_file}")
                            f.create_dataset('defect_tiff', data=defect_tiff)

                    elif 'Defect_' + str(k) + '_Reference_1_Topography3' in img:
                        if 'reference_tiff' not in f.keys():
                            reference_tiff = np.array(Image.open(img))
                            print(f"add reference_file from {img} to {dst_h5_file}")
                            f.create_dataset('reference_tiff', data=reference_tiff)
                f.close()
                f_read = h5py.File(dst_h5_file, 'r')
                if 'class_tiff' in f_read.keys() and 'defect_tiff' in f_read.keys() and 'reference_tiff' in f_read.keys():
                    f_read.close()
                    print(f"======================{dst_h5_file_name} finished====klarf : {file}==================")
                else:
                    f_read.close()
                    os.remove(dst_h5_file)
                    print(f"xxxxxxxxxxxxxxxxxxxxxxxxxxxx{dst_h5_file_name} error xxxxxxxxxxxxxxxxxxxxxxxxxx")




def getClassImgs():
    dir = r'C:\Users\caoxijian\Desktop\origin_datasets\klarf\xy'
    dst_dir = r'D:\proj\data_xy'

    if not os.path.exists(dst_dir):
        os.mkdir(dst_dir)

    for category in category2id.keys():
        category_dir = os.path.join(dst_dir, category)
        if not os.path.exists(category_dir):
            os.mkdir(category_dir)

    klarf_files = os.listdir(dir)
    for file in klarf_files:
        filePath = os.path.join(dir, file)
        deviceId, stepId, lotId, waferId, defectDict = readKlarf(filePath)
        src_dir = os.path.join('Z:\\' + deviceId, stepId, lotId, waferId)
        for k, v in defectDict.items(): # k: defect_id   v: defect_type
            if str(v) in category2id.values():
                img_dir = src_dir + r'\*\Images\Run_1\Defect_' + str(k) + '_Class_1_Topography3.tiff'
                img_list = glob.glob((img_dir))
                #['Z:\\Y-MS42\\AL_ASI_MINI\\lDE1621.04_0718\\w21\\BB94982A-9250-4CF7-9B04-71E82FCB2B87\\Images\\Run_1\\Defect_99982_Class_1_Topography3.tiff']
                for img in img_list:
                    category = id2category[str(v)]
                    # copy files
                    dst_filename = category + '_' + deviceId + '_' + stepId + '_' + lotId + '_' + waferId + '_' + str(k) + '_' + 'Defect_Topography3.tiff'
                    dst_file = os.path.join(dst_dir, category, dst_filename)
                    print(f"copy files from {img} to {dst_file}")
                    shutil.copyfile(img, dst_file)

def getDefectImgs():
    dir = r'C:\Users\caoxijian\Desktop\origin_datasets\klarf\total04_10'
    dst_dir = r'D:\proj\dataset_h5'
    if not os.path.exists(dst_dir):
        os.mkdir(dst_dir)

    for category in category2id.keys():
        category_dir = os.path.join(dst_dir, category)
        if not os.path.exists(category_dir):
            os.mkdir(category_dir)

    klarf_files = os.listdir(dir)
    for file in klarf_files:
        filePath = os.path.join(dir, file)
        deviceId, stepId, lotId, waferId, defectDict = readKlarf(filePath)
        src_dir = os.path.join('Z:\\' + deviceId, stepId, lotId, waferId)
        for k, v in defectDict.items(): # k: defect_id   v: defect_type
            if str(v) in category2id.values():
                img_dir = src_dir + r'\*\Images\Run_1\Defect_' + str(k) + '_Defect_1_Topography3.tiff'
                img_list = glob.glob((img_dir))
                #['Z:\\Y-MS42\\AL_ASI_MINI\\lDE1621.04_0718\\w21\\BB94982A-9250-4CF7-9B04-71E82FCB2B87\\Images\\Run_1\\Defect_99982_Class_1_Topography3.tiff']
                for img in img_list:
                    category = id2category[str(v)]
                    # copy files
                    dst_filename = category + '_' + deviceId + '_' + stepId + '_' + lotId + '_' + waferId + '_' + str(k) + '_' + 'Defect_Topography3.tiff'
                    dst_file = os.path.join(dst_dir, category, dst_filename)
                    print(f"copy files from {img} to {dst_file}")
                    shutil.copyfile(img, dst_file)

def getReferenceImgs():
    dir = r'C:\Users\caoxijian\Desktop\origin_datasets\klarf\total04_10'
    dst_dir = r'D:\proj\defect_class_images_55'

    if not os.path.exists(dst_dir):
        os.mkdir(dst_dir)

    for category in category2id.keys():
        category_dir = os.path.join(dst_dir, category)
        if not os.path.exists(category_dir):
            os.mkdir(category_dir)

    klarf_files = os.listdir(dir)
    for file in klarf_files:
        filePath = os.path.join(dir, file)
        deviceId, stepId, lotId, waferId, defectDict = readKlarf(filePath)
        src_dir = os.path.join('Z:\\' + deviceId, stepId, lotId, waferId)
        for k, v in defectDict.items(): # k: defect_id   v: defect_type
            if str(v) in category2id.values():
                img_dir = src_dir + r'\*\Images\Run_1\Defect_' + str(k) + '_Defect_1_Topography3.tiff'
                img_list = glob.glob((img_dir))
                #['Z:\\Y-MS42\\AL_ASI_MINI\\lDE1621.04_0718\\w21\\BB94982A-9250-4CF7-9B04-71E82FCB2B87\\Images\\Run_1\\Defect_99982_Class_1_Topography3.tiff']
                for img in img_list:
                    category = id2category[str(v)]
                    # copy files
                    dst_filename = category + '_' + deviceId + '_' + stepId + '_' + lotId + '_' + waferId + '_' + str(k) + '_' + 'Defect_Topography3.tiff'
                    dst_file = os.path.join(dst_dir, category, dst_filename)
                    print(f"copy files from {img} to {dst_file}")
                    shutil.copyfile(img, dst_file)


def check_dir():
    dir = r'D:\proj\defect_class_images_55'
    categories = os.listdir(dir)
    for category in categories:
        category_dir = os.path.join(dir, category)
        print(f"len of {category}: {len(os.listdir(category_dir))}")

def sampleH5():
    #已sample：Reference、Burried_particle、Surface_particle、CMP_scratch、Cu_diffuse、Prelayer_defect_developed、Residue、Void
    dir = r'D:\proj\dataset_h5\Void'
    sample_dir = r'D:\proj\dataset_h5_sample\Void'
    h5total = os.listdir(dir)
    np.random.shuffle(h5total)
    h5sample = h5total[0: 6000]
    for h5 in h5sample:
        src_file = os.path.join(dir, h5)
        dst_file = os.path.join(sample_dir, h5)
        shutil.copyfile(src_file, dst_file)





def sample():
    dir = r'D:\proj\defect_class_images_2024'
    sample_dir = r'D:\proj\defect_class_images_2024_sampled'
    if not os.path.exists(sample_dir):
        os.mkdir(sample_dir)
    categories = os.listdir(dir)
    for category in categories:
        category_dir = os.path.join(dir, category)
        category_sample_dir = os.path.join(sample_dir, category)
        if not os.path.exists(category_sample_dir):
            os.mkdir(category_sample_dir)

        category_num = len(os.listdir(category_dir))
        category_imgs = os.listdir(category_dir)
        np.random.shuffle(category_imgs)

        if category_num > 5000:
            category_imgs = category_imgs[0: 5000]


        for img in category_imgs:
            srcimg = os.path.join(category_dir, img)
            dstimg = os.path.join(category_sample_dir, img)
            shutil.copyfile(srcimg, dstimg)

    def sample2():
        dir = r'C:\Users\caoxijian\Desktop\origin_datasets\class_images'
        sample_dir = r'C:\Users\caoxijian\Desktop\origin_datasets\class_images_sampled'
        if not os.path.exists(sample_dir):
            os.mkdir(sample_dir)
        categories = os.listdir(dir)
        for category in categories:
            category_dir = os.path.join(dir, category)
            category_sample_dir = os.path.join(sample_dir, category)
            if not os.path.exists(category_sample_dir):
                os.mkdir(category_sample_dir)

            category_num = len(os.listdir(category_dir))
            category_imgs = os.listdir(category_dir)
            np.random.shuffle(category_imgs)

            if category_num > 5000:
                category_imgs = category_imgs[0: 3000]
            elif category_num > 3000:
                category_imgs = category_imgs[0: 2500]

            for img in category_imgs:
                srcimg = os.path.join(category_dir, img)
                dstimg = os.path.join(category_sample_dir, img)
                shutil.copyfile(srcimg, dstimg)


def sample_test():
    dir = r'C:\Users\caoxijian\Desktop\origin_datasets\example_images'
    sample_dir = r'C:\Users\caoxijian\Desktop\origin_datasets\example_images_sampled'
    if not os.path.exists(sample_dir):
        os.mkdir(sample_dir)
    categories = os.listdir(dir)
    for category in categories:
        category_dir = os.path.join(dir, category)
        category_sample_dir = os.path.join(sample_dir, category)
        if not os.path.exists(category_sample_dir):
            os.mkdir(category_sample_dir)

        category_num = len(os.listdir(category_dir))
        category_imgs = os.listdir(category_dir)
        np.random.shuffle(category_imgs)

        if category_num > 200:
            category_imgs = category_imgs[0: 200]

        for img in category_imgs:
            srcimg = os.path.join(category_dir, img)
            dstimg = os.path.join(category_sample_dir, img)
            shutil.copyfile(srcimg, dstimg)


if __name__ == '__main__':
    # sampleFromklarf()
    check_data()
    # getClassImgs()
    # check_dir()
    # sample()
    # sample_test()
    # getH5Data()
    # sampleH5()

'''
0821:6498
0814:5899
0807:5294
0718:6792
0710:7202
0703:6000
0607:4700
0529:2495
0524:3698
0515:2599

total:51177
'''

'''
[(1, 26211), (7, 6166), (6, 2962), (4, 2727), (107, 2531), (11, 2519), (3, 1808), (5, 1113), (67, 982), (2, 842), 
(14, 685), (84, 525), (83, 499), (89, 463), (74, 411), (174, 149), (76, 134), (88, 100), (0, 84), (12, 59), (65, 53), 
(62, 50), (13, 27), (24, 21), (68, 15), (52, 12), (148, 9), (28, 8), (214, 7), (167, 1), (183, 1), (86, 1), (19, 1), (26, 1)]

'''

'''
len of ADI_particle_deve: 3448
len of ArrayPeeling: 427
len of Bump: 4135
len of Burried_particle: 5457
len of CMP_scratch: 7003
len of Cu_diffuse: 3266
len of Cu_loss: 3352
len of OtherPeeling: 961
len of Oxide_scratch: 4542
len of Pattern_fail: 2104
len of Pre_layer_def: 3464
len of Reference: 25264
len of Residue: 12802
len of Surface_particle: 9141
len of Void: 8563
'''

'''
[(1, 135529), (7, 27233), (107, 18341), (11, 18062), (6, 14278), (4, 10159), (3, 6588), (5, 6302), (14, 5423), (84, 4409), 
(83, 4345), (67, 4276), (2, 3291), (74, 3117), (89, 1935), (0, 656), (76, 588), (88, 523), (24, 453), (174, 406), 
(13, 384), (35, 306), (52, 290), (62, 264), (71, 208), (12, 203), (21, 175), (27, 119), (148, 117), (33, 108), 
(8, 88), (65, 82), (147, 75), (19, 73), (68, 47), (214, 46), (114, 40), (189, 34), (16, 32), (28, 20), 
(117, 11), (75, 10), (54, 9), (26, 7), (211, 4), (41, 4), (47, 3), (124, 3), (188, 2), (79, 1), 
(167, 1), (183, 1), (86, 1), (133, 1), (235, 1)]
<class 'list'>
55
'''